#!/bin/bash

conda_create_env() {
  local exit_code=0

  if [ -n "$GITHUB_PROXY" ]; then
    # clear default channels
    export CONDA_DEFAULT_CHANNELS=""
    conda config --remove channels defaults 2>/dev/null
    conda config --system --remove channels defaults 2>/dev/null
  fi

  export CONDA_PKGS_DIRS="${LMD_DATA_ROOT}/install-cache/conda"
  conda create -p "$@"
  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "conda cache broken. Cleaning cache..."
    conda clean --all -y
  fi

  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "conda create again 1"
    conda create -p "$@"
  fi

  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "conda cache broken. Cleaning cache..."
    conda clean --all -y
  fi

  # clean conda channels
  if [ $exit_code -ne 0 ]; then
    TMP_CONDA_CHANNELS=$CONDA_CHANNELS
    export CONDA_CHANNELS=$CONDA_CHANNELS_SUB

    conda config --set ssl_verify_peer false
    conda config --set ssl_verify false

    echo "exit_code: "$exit_code
    echo "conda create again 2 with CHANNELS_SUB ${CONDA_CHANNELS} and disable ssl verify."
    conda create -p "$@"

    # export CONDA_CHANNELS=$TMP_CONDA_CHANNELS
  fi

  exit_code=$?
  if [ $exit_code -eq 1 ]; then
    echo "conda cache broken. Cleaning cache..."
    conda clean --all -y
  fi

  # restore default channels
  unset CONDA_DEFAULT_CHANNELS
  conda config --add channels defaults
  conda config --system --add channels defaults

  # clean conda channels. use default channels.
  if [ $exit_code -ne 0 ]; then    
    TMP_CONDA_CHANNELS=$CONDA_CHANNELS
    export CONDA_CHANNELS=""
    echo "exit_code: "$exit_code
    echo "conda create again 3"
    conda create -c defaults -p "$@"

    export CONDA_CHANNELS=$TMP_CONDA_CHANNELS
  fi
}

conda_run_cmd() {
  local exit_code=0

  export CONDA_PKGS_DIRS="${LMD_DATA_ROOT}/install-cache/conda"
  conda run -p "$CONDA_PREFIX" --no-capture-output "$@"
  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "conda cache broken. Cleaning cache..."
    conda clean --all -y
  fi

  if [ $exit_code -ne 0 ]; then
    echo "conda run again"
    echo "exit_code: "$exit_code
    conda run -p "$CONDA_PREFIX" --no-capture-output "$@"
  fi

}

exchange_pip_index() {
  TMP_PIP_INDEX_URL=$PIP_INDEX_URL
  TMP_PIP_INDEX_URL_SUB=$PIP_INDEX_URL_SUB

  export PIP_INDEX_URL=$TMP_PIP_INDEX_URL_SUB
  export PIP_INDEX_URL_SUB=$TMP_PIP_INDEX_URL  
  echo "exchange_pip_index" 
  echo "index: $PIP_INDEX_URL" 
  echo "index_sub: $PIP_INDEX_URL_SUB" 
}


hide_pip_index() {
  HIDE_PIP_INDEX_URL=$PIP_INDEX_URL
  HIDE_PIP_INDEX_URL_SUB=$PIP_INDEX_URL_SUB

  export PIP_INDEX_URL=""
  export PIP_INDEX_URL_SUB=""
  echo "hide_pip_index"
}

restore_pip_index() {
  export PIP_INDEX_URL="$HIDE_PIP_INDEX_URL"
  export PIP_INDEX_URL_SUB="$HIDE_PIP_INDEX_URL_SUB"
  echo "restore_pip_index"
}

conda_run_pip_install() {
  local exit_code=0
  
  export CONDA_PKGS_DIRS="${LMD_DATA_ROOT}/install-cache/conda"
  PIP_CACHE_DIR="${LMD_DATA_ROOT}/install-cache/pip"

  echo "PIP_CACHE_DIR $PIP_CACHE_DIR"

  conda run -p "$CONDA_PREFIX" --no-capture-output pip install --cache-dir="$PIP_CACHE_DIR" "$@"
  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "Cleaning conda cache..."
    conda clean --all -y
  fi

  if [ $exit_code -ne 0 ]; then
    echo "pip install again"
    echo "exit_code: "$exit_code
    exchange_pip_index
    conda run -p "$CONDA_PREFIX" --no-capture-output pip install --cache-dir="$PIP_CACHE_DIR" "$@"
  fi

  if [ $exit_code -ne 0 ]; then
    echo "pip install again2"
    exchange_pip_index
    conda run -p "$CONDA_PREFIX" --no-capture-output pip install --cache-dir="$PIP_CACHE_DIR" "$@"
    # restore_pip_index
  fi
}

conda_run_hf_download() {
  local exit_code=0
  echo "hf download use endpoint: $HF_ENDPOINT"

  conda run -p "$CONDA_PREFIX" --no-capture-output hf download "$@"

  exit_code=$?
  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "run hf download again 1"
    conda run -p "$CONDA_PREFIX" --no-capture-output hf download "$@"
  fi

  exit_code=$?

  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "run hf download again 2. disable hf xet"
    export HF_HUB_DISABLE_XET=1
    conda run -p "$CONDA_PREFIX" --no-capture-output hf download "$@"

    echo "enable hf xet"
    export HF_HUB_DISABLE_XET=0
  fi

  exit_code=$?

  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "run hf download again 3. disable hf xet"
    export HF_HUB_DISABLE_XET=1
    conda run -p "$CONDA_PREFIX" --no-capture-output hf download "$@"

    echo "enable hf xet"
    export HF_HUB_DISABLE_XET=0
  fi
}

# use modelscope to downlaod models and change dir structure like hf
conda_run_hf_dl_by_modelsople() {
  local modelId=$1
  local localModelDir=$2
  local transformersCacheDir=$HF_HOME
  
  local exit_code=0
  echo "modelscope download: $modelId to $localModelDir"

  local fullLocalModelDir="${transformersCacheDir}/hub/models--${localModelDir}/snapshots/latest"
  mkdir -p $fullLocalModelDir

  local fullLocallRefsDir="${transformersCacheDir}/hub/models--${localModelDir}/refs"
  mkdir -p "$fullLocallRefsDir"
  cp "$LMD_BASE_INSTALL_SCRIPT_DIR/global/main" "${fullLocallRefsDir}/main"
 
  conda run -p "$CONDA_PREFIX" --no-capture-output modelscope download $modelId --local_dir "$fullLocalModelDir"

  exit_code=$?
  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "modelscope hf download again"
    conda run -p "$CONDA_PREFIX" --no-capture-output modelscope download $modelId --local_dir "$fullLocalModelDir"
  fi
}

conda_run_conda_forge_install() {
  local exit_code=0

  export CONDA_PKGS_DIRS="${LMD_DATA_ROOT}/install-cache/conda"
  echo "CONDA_PKGS_DIRS $CONDA_PKGS_DIRS"

  if [ -n "$GITHUB_PROXY" ]; then
    # clear default channels
    export CONDA_DEFAULT_CHANNELS=""
    conda config --remove channels defaults 2>/dev/null
    conda config --system --remove channels defaults 2>/dev/null
  fi

  if [ -n "$GITHUB_PROXY" ]; then
    CONDA_FORGE_VALUE="https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge"
  else 
    CONDA_FORGE_VALUE="conda-forge"
  fi

  # first time
  conda run -p "$CONDA_PREFIX" --no-capture-output conda install -y -c $CONDA_FORGE_VALUE "$@"

  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "Cleaning conda cache..."
    conda clean --all -y
  fi

  if [ -n "$GITHUB_PROXY" ]; then
    CONDA_FORGE_VALUE="https://mirror.nju.edu.cn/anaconda/cloud/conda-forge"
  else 
    CONDA_FORGE_VALUE="conda-forge"
  fi

  # second time
  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    conda run -p "$CONDA_PREFIX" --no-capture-output conda install -y -c $CONDA_FORGE_VALUE "$@"
  fi

  exit_code=$?

  if [ $exit_code -eq 1 ]; then
    echo "Cleaning conda cache..."
    conda clean --all -y
  fi

  # restore default channels
  unset CONDA_DEFAULT_CHANNELS
  conda config --add channels defaults
  conda config --system --add channels defaults
  # third time
  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "use default conda-forge"
    conda run -p "$CONDA_PREFIX" --no-capture-output conda install -y -c conda-forge "$@"
  fi
}


conda_run_modelsople_download() {
  local modelId=$1
  local localModelDir=$2
  
  local exit_code=0
  echo "modelscope download: $modelId  files: $localModelDir"
 
  conda run -p "$CONDA_PREFIX" --no-capture-output modelscope download $modelId --local_dir "$localModelDir"

  exit_code=$?
  if [ $exit_code -ne 0 ]; then
    echo "exit_code: "$exit_code
    echo "modelscope download again"
    conda run -p "$CONDA_PREFIX" --no-capture-output modelscope download $modelId --local_dir "$localModelDir"
  fi
}
